package main

import (
    "context"
    "fmt"
    "log"
    "math"
    "math/rand"
    "sort"
    "strconv"
    "sync"
    "testing"
    "time"

    "github.com/milvus-io/milvus/client/v2/milvusclient"
    "github.com/milvus-io/milvus/client/v2/entity"
    "github.com/milvus-io/milvus/client/v2/index"
    "github.com/milvus-io/milvus/client/v2/column"
)

// 编译
// go test -c ascend_interface_test.go -o ascend_interface_test
// 运行
// ./ascend_interface_test --test.run=TestXXXXXX --test.v

var (
    nEntities, dim = 3000, 512
)

const (
    milvusAddr     = `172.18.0.4:19666`
    // milvusAddr     = `localhost:19530`
    collectionName = "hello_ascend_test"

    msgFmt                         = "\n==== %s ====\n"
    idCol, randomCol, embeddingCol = "ID", "random", "embeddings"
    topK                           = 10
    maxLength                      = 10
    insertLimit                    = 10000
)

func Min(x, y int) int {
    if x < y {
        return x
    }
    return y
}

func Shuffle(num int) []int {
    var slice []int
    for i := 0; i < num; i++ {
        slice = append(slice, i)
    }
    r := rand.New(rand.NewSource(time.Now().Unix()))
    for len(slice) > 0 {
        n := len(slice)
        randIndex := r.Intn(n)
        slice[n-1], slice[randIndex] = slice[randIndex], slice[n-1]
        slice = slice[:n-1]
    }
    return slice
}

// 创建collection，创建index，insert数据，flush，search query，对比结果
func TestSearchAndQueryVarCharID(t *testing.T) {
    ctx := context.Background()

    fmt.Printf(msgFmt, "start connecting to Milvus")
    c, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
        Address: milvusAddr,
    })
    if err != nil {
        log.Fatalf("failed to connect to milvus, err: %v", err)
    }
    defer c.Close(ctx)

    // delete collection if exists
    has, err := c.HasCollection(ctx, milvusclient.NewHasCollectionOption(collectionName))
    if err != nil {
        log.Fatalf("failed to check collection exists, err: %v", err)
    }
    if has {
        fmt.Printf(msgFmt, "collection exist, droping...")
        c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
    }

    // create collection
    fmt.Printf(msgFmt, "create collection `hello_ascend_test")
    schema := &entity.Schema{
        CollectionName: collectionName,
        Description:    "hello_ascend_test is the simplest demo to introduce the APIs",
        AutoID:         false,
        Fields: []*entity.Field{
            {
                Name:       idCol,
                DataType:   entity.FieldTypeVarChar,
                PrimaryKey: true,
                AutoID:     false,
                TypeParams: map[string]string{
                    entity.TypeParamMaxLength: fmt.Sprintf("%d", maxLength),
                },
            },
            {
                Name:       randomCol,
                DataType:   entity.FieldTypeDouble,
                PrimaryKey: false,
                AutoID:     false,
            },
            {
                Name:       embeddingCol,
                DataType:   entity.FieldTypeFloatVector,
                TypeParams: map[string]string{
                    entity.TypeParamDim: fmt.Sprintf("%d", dim),
                },
            },
        },
    }

    if err := c.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
        WithShardNum(1)); err != nil {
        log.Fatalf("create collection failed, err: %v", err)
    }

    // build index
    fmt.Printf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    idx := index.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        log.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v", err)
    }
    indexTask, err := c.CreateIndex(ctx, milvusclient.NewCreateIndexOption(collectionName, embeddingCol, idx))
    if err != nil {
        log.Fatalf("failed to create index, err: %v", err)
    }
    if err = indexTask.Await(ctx); err != nil {
        log.Fatalf("Index building failed: %v", err)
    }

    // insert data
    fmt.Printf(msgFmt, "start inserting random entities")
    idList := make([]string, 0, nEntities)
    randomList := make([]float64, 0, nEntities)
    embeddingList := make([][]float32, 0, nEntities)

    rand.Seed(time.Now().UnixNano())

    // generate data
    for i := 0; i < nEntities; i++ {
        idList = append(idList, fmt.Sprintf("%d", i))
        randomList = append(randomList, rand.Float64())
        vec := make([]float32, 0, dim)
        var sum float32 = 0.0
        for j := 0; j < dim; j++ {
            vec = append(vec, rand.Float32())
            sum += vec[j] * vec[j]
        }
        sum = float32(math.Sqrt(float64(sum)))
        for j := 0; j < dim; j++ {
            vec[j] = vec[j] / sum
        }
        embeddingList = append(embeddingList, vec)
    }

    idColData := column.NewColumnVarChar(idCol, idList)
    randomColData := column.NewColumnDouble(randomCol, randomList)
    embeddingColData := column.NewColumnFloatVector(embeddingCol, dim, embeddingList)

    opt := milvusclient.NewColumnBasedInsertOption(
        collectionName,
        idColData,
        randomColData,
        embeddingColData,
    )
    if _, err := c.Insert(ctx, opt); err != nil {
        log.Fatalf("failed to insert random data into `hello_ascend_test`, err: %v", err)
    }

    flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        log.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        log.Fatalf("flush failed: %v", err)
    }

    fmt.Printf(msgFmt, "start loading collection")
    loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        log.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        log.Fatalf("loading collection failed: %v", err)
    }

    fmt.Printf(msgFmt, "start searching based on vector similarity")
    vec2search := []entity.Vector{
        entity.FloatVector(embeddingList[len(embeddingList)-2]),
        entity.FloatVector(embeddingList[len(embeddingList)-1]),
    }
    begin := time.Now()
    sRet, err := c.Search(ctx, milvusclient.NewSearchOption(collectionName, topK, vec2search).
        WithOutputFields(randomCol).WithANNSField(embeddingCol))
    end := time.Now()
    if err != nil {
        log.Fatalf("failed to search collection, err: %v", err)
    }

    fmt.Println("results:")
    for _, res := range sRet {
        printResult(&res)
    }
    fmt.Printf("\tsearch latency: %dms\n", end.Sub(begin)/time.Millisecond)

    // hybrid search
    fmt.Printf(msgFmt, "start hybrid searching with `random > 0.5`")
    begin = time.Now()
    sRet2, err := c.Search(ctx, milvusclient.NewSearchOption(collectionName, topK, vec2search).
        WithFilter("random > 0.5").WithOutputFields(randomCol).WithANNSField(embeddingCol))
    end = time.Now()
    if err != nil {
        log.Fatalf("failed to search collection, err: %v", err)
    }
    fmt.Println("results:")
    for _, res := range sRet2 {
        printResult(&res)
    }
    fmt.Printf("\tsearch latency: %dms\n", end.Sub(begin)/time.Millisecond)

    // delete data
    fmt.Printf(msgFmt, "start deleting with expr ``")
    pks := column.NewColumnVarChar(idCol, []string{"0", "1"})
    sRet3, err := c.Get(ctx, milvusclient.NewQueryOption(collectionName).
        WithOutputFields(randomCol).
        WithIDs(pks))
    if err != nil {
        log.Fatalf("failed to query result, err: %v", err)
    }
    fmt.Println("results:")
    idList = make([]string, 0)
    randList := make([]float64, 0)

    for _, col := range sRet3.Fields {
        if col.Name() == idCol {
            idColumn := col.(*column.ColumnVarChar)
            for i := 0; i < col.Len(); i++ {
                val, err := idColumn.Value(i)
                if err != nil {
                    log.Fatal(err)
                }
                idList = append(idList, val)
            }
        } else {
            randColumn := col.(*column.ColumnDouble)
            for i := 0; i < col.Len(); i++ {
                val, err := randColumn.Value(i)
                if err != nil {
                    log.Fatal(err)
                }
                randList = append(randList, val)
            }
        }
    }
    fmt.Printf("\tids: %#v, randoms: %#v\n", idList, randList)

    _, err = c.Delete(ctx, milvusclient.NewDeleteOption(collectionName).
        WithStringIDs(idCol, []string{"0", "1"}))
    if err != nil {
        log.Fatalf("failed to delete by pks, err: %v", err)
    }
    _, err = c.Get(ctx, milvusclient.NewQueryOption(collectionName).
        WithOutputFields(randomCol).
        WithIDs(pks))
    if err != nil {
        log.Printf("failed to query result, err: %v", err)
    }

    // drop collection
    fmt.Printf(msgFmt, "drop collection `hello_ascend_test`")
    if err := c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName)); err != nil {
        log.Fatalf("failed to drop collection, err: %v", err)
    }
}

// 创建collection，创建index，insert数据，flush，search query，对比结果
func TestSearchAndQueryIntID(t *testing.T) {
    ctx := context.Background()
    t.Logf(msgFmt, "start connecting to Milvus")
    c, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
        Address: milvusAddr,
    })
    if err != nil {
        t.Fatalf("failed to connect to milvus, err: %v", err)
    }
    defer c.Close(ctx)

    // delete collection if exists
    has, err := c.HasCollection(ctx, milvusclient.NewHasCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to check collection exists, err: %v", err)
    }
    if has {
        c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
    }

    // create collection
    t.Logf(msgFmt, "create collection `hello_ascend_test`")
    schema := CreateSchema()

    if err := c.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
                                 WithShardNum(1)); err != nil {
        t.Fatalf("create collection failed, err: %v", err)
    }
    defer func() {
        t.Logf(msgFmt, "drop collection")
        if err := c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName)); err != nil {
            t.Fatalf("failed to drop collection, err: %v", err)
        }
    }()

    // build index
    t.Logf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    idx := index.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        t.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v", err)
    }
    indexTask, err := c.CreateIndex(ctx, milvusclient.NewCreateIndexOption(collectionName, embeddingCol, idx))
    if err != nil {
        t.Fatalf("failed to create index, err: %v", err)
    }
    if err = indexTask.Await(ctx); err != nil {
        t.Fatalf("Index building failed: %v", err)
    }

    // generate data
    idList, randomList, embeddingList := GenerateData(t, nEntities, 0)

    // insert data
    InsertPaged(t, ctx, c, idList, randomList, embeddingList, "")
    flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }

    // load collection
    t.Logf(msgFmt, "start loading collection")
    loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        t.Fatalf("loading collection failed: %v", err)
    }

    Search(t, ctx, c, idList, randomList, embeddingList, []string{})
    Query(t, ctx, c, randomList)

    SearchExpr(t, ctx, c, embeddingList, []string{})

    if stats, err := c.GetCollectionStats(ctx, milvusclient.NewGetCollectionStatsOption(collectionName)); err != nil {
        t.Fatalf("failed to get collection statistics, err: %v", err)
    } else {
        for key, val := range stats {
            t.Logf("collection %s [%s : %s]", collectionName, key, val)
        }
    }
}

// 创建3个Partition，分别插入不同的数据，分别search
func TestPartition(t *testing.T) {
    ctx := context.Background()
    t.Logf(msgFmt, "start connecting to Milvus")
    c, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
        Address: milvusAddr,
    })
    if err != nil {
        t.Fatalf("failed to connect to milvus, err: %v", err)
    }
    defer c.Close(ctx)

    // delete collection if exists
    has, err := c.HasCollection(ctx, milvusclient.NewHasCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to check collection exists, err: %v", err)
    }
    if has {
        c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
    }

    // create collection
    t.Logf(msgFmt, "create collection `hello_ascend_test`")
    schema := CreateSchema()
    if err := c.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
                                 WithShardNum(1)); err != nil {
        t.Fatalf("create collection failed, err: %v", err)
    }
    defer func() {
        t.Logf(msgFmt, "drop collection")
        if err := c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName)); err != nil {
            t.Fatalf("failed to drop collection, err: %v", err)
        }
    }()

    // build index
    t.Logf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    idx := index.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        t.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v\n", err)
    }
    indexTask, err := c.CreateIndex(ctx, milvusclient.NewCreateIndexOption(collectionName, embeddingCol, idx))
    if err != nil {
        t.Fatalf("failed to create index, err: %v", err)
    }
    if err = indexTask.Await(ctx); err != nil {
        t.Fatalf("Index building failed: %v", err)
    }

    t.Logf(msgFmt, "list collection `hello_ascend_test`")
    collections, err := c.ListCollections(ctx, milvusclient.NewListCollectionOption())
    if err != nil {
        log.Fatal("failed to list collections:", err.Error())
    }
    t.Logf("Collection name: %v\n", collections)

    t.Logf(msgFmt, "show collection partitions")
    partitions, err := c.ListPartitions(ctx, milvusclient.NewListPartitionOption(collectionName))
    if err != nil {
        t.Fatal("failed to show partition:", err.Error())
    }
    t.Logf("partitions: %v\n", partitions)

    t.Logf(msgFmt, "create partitions")
    partitionNames := []string{"partition_1", "partition_2", "partition_3"}
    for _, name := range partitionNames {
        if err = c.CreatePartition(ctx, milvusclient.NewCreatePartitionOption(collectionName, name)); err != nil {
            t.Fatalf("failed to create partition %s, %v\n", name, err)
        }
        has, err = c.HasPartition(ctx, milvusclient.NewHasPartitionOption(collectionName, name))
        if err != nil {
            t.Fatalf("failed to check partition %s, %v\n", name, err)
        }
        if !has {
            t.Fatalf("partition %s created, but not exist", name)
        }
    }

    t.Logf(msgFmt, "show collection partitions after create")
    // show collection partitions, check creation
    partitions, err = c.ListPartitions(ctx, milvusclient.NewListPartitionOption(collectionName))
    if err != nil {
        t.Fatal("failed to show partitions:", err.Error())
    }
    t.Logf("partition name: %v\n", partitions)

    // generate data
    t.Logf(msgFmt, "start generating data")
    idList, randomList, embeddingList := GenerateData(t, nEntities*len(partitionNames), 0)
    // insert data
    for i, name := range partitionNames {
        InsertPaged(t, ctx, c, idList[i*nEntities:(i+1)*nEntities], randomList[i*nEntities:(i+1)*nEntities],
            embeddingList[i*nEntities:(i+1)*nEntities], name)
    }
    flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }

    // load collection
    t.Logf(msgFmt, "start loading collection")
    loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        t.Fatalf("loading collection failed: %v", err)
    }

    t.Logf(msgFmt, "start searching partition")
    for i, name := range partitionNames {
        Search(t, ctx, c, idList[i*nEntities:(i+1)*nEntities], randomList[i*nEntities:(i+1)*nEntities],
            embeddingList[i*nEntities:(i+1)*nEntities], []string{name})
    }

    t.Logf(msgFmt, "start release partition")
    for _, name := range partitionNames {
        if err := c.ReleasePartitions(ctx, milvusclient.NewReleasePartitionsOptions(collectionName, name)); err != nil {
            t.Fatal("failed to release partitions:", err.Error())
        }
    }

    t.Logf(msgFmt, "start show partition")
    partitions, err = c.ListPartitions(ctx, milvusclient.NewListPartitionOption(collectionName))
    if err != nil {
        t.Fatal("failed to show partitions:", err.Error())
    }
    t.Logf("partition name: %v\n", partitions)

    t.Logf(msgFmt, "start generating new data")
    // generate new data
    idList, randomList, embeddingList = GenerateData(t, nEntities*len(partitionNames), nEntities*len(partitionNames))
    // insert data
    for i, name := range partitionNames {
        InsertPaged(t, ctx, c, idList[i*nEntities:(i+1)*nEntities], randomList[i*nEntities:(i+1)*nEntities],
                    embeddingList[i*nEntities:(i+1)*nEntities], name)
    }
    flushtask, err = c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }
    // load collection
    t.Logf(msgFmt, "start loading collection")
    loadTask, err = c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        t.Fatalf("loading collection failed: %v", err)
    }
    t.Logf(msgFmt, "start searching new partition")
    for i, name := range partitionNames {
        Search(t, ctx, c, idList[i*nEntities:(i+1)*nEntities], randomList[i*nEntities:(i+1)*nEntities],
            embeddingList[i*nEntities:(i+1)*nEntities], []string{name})
    }

    t.Logf(msgFmt, "drop collection partitions")
    for _, name := range partitionNames {
        if err = c.ReleasePartitions(ctx, milvusclient.NewReleasePartitionsOptions(collectionName, name)); err != nil {
            t.Fatalf("failed to release partition %s, %v\n", name, err)
        }
        if err = c.DropPartition(ctx, milvusclient.NewDropPartitionOption(collectionName, name)); err != nil {
            t.Fatalf("failed to drop partition %s, %v\n", name, err)
        }
        has, err = c.HasPartition(ctx, milvusclient.NewHasPartitionOption(collectionName, name))
        if err != nil {
            t.Fatalf("failed to check partition %s, %v\n", name, err)
        }
        if has {
            t.Fatalf("partition %s dropped, but exist\n", name)
        }
    }
    if stats, err := c.GetCollectionStats(ctx, milvusclient.NewGetCollectionStatsOption(collectionName)); err != nil {
        t.Fatalf("failed to get collection statistics, err: %v", err)
    } else {
        for key, val := range stats {
            t.Logf("collection %s [%s : %s]", collectionName, key, val)
        }
    }
}

// 创建collection，insert数据，flush，search；删除DeleteByPks，再Search
func TestDeleteSearch(t *testing.T) {
    ctx := context.Background()
    t.Logf(msgFmt, "start connecting to Milvus")
    c, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
        Address: milvusAddr,
    })
    if err != nil {
        t.Fatalf("failed to connect to milvus, err: %v\n", err)
    }
    defer c.Close(ctx)

    // delete collection if exists
    has, err := c.HasCollection(ctx, milvusclient.NewHasCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to check collection exists, err: %v\n", err)
    }
    if has {
        c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
    }

    // create collection
    t.Logf(msgFmt, "create collection `hello_ascend_test`")
    schema := CreateSchema()

    if err := c.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
        WithShardNum(1)); err != nil {
        t.Fatalf("create collection failed, err: %v\n", err)
    }
    defer func() {
        t.Logf(msgFmt, "drop collection")
        if err := c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName)); err != nil {
            t.Fatalf("failed to drop collection, err: %v\n", err)
        }
    }()

    // build index
    t.Logf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    idx := index.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        t.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v\n", err)
    }
    indexTask, err := c.CreateIndex(ctx, milvusclient.NewCreateIndexOption(collectionName, embeddingCol, idx))
    if err != nil {
        t.Fatalf("failed to create index, err: %v", err)
    }
    if err = indexTask.Await(ctx); err != nil {
        t.Fatalf("Index building failed: %v", err)
    }

    partitionName := "new_partition_1"
    if err = c.CreatePartition(ctx, milvusclient.NewCreatePartitionOption(collectionName, partitionName)); err != nil {
        t.Fatalf("failed to create partition %s, %v\n", partitionName, err)
    }
    has, err = c.HasPartition(ctx, milvusclient.NewHasPartitionOption(collectionName, partitionName))
    if err != nil {
        t.Fatalf("failed to check partition %s, %v\n", partitionName, err)
    }
    if !has {
        t.Fatalf("partition %s created, but not exist", partitionName)
    }

    // generate data
    idList, randomList, embeddingList := GenerateData(t, nEntities, 0)

    // insert data
    InsertPaged(t, ctx, c, idList, randomList, embeddingList, partitionName)
    flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }

    // load collection
    t.Logf(msgFmt, "start loading collection")
    loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        t.Fatalf("loading collection failed: %v", err)
    }

    t.Logf(msgFmt, "start searching")
    var vec2search []entity.Vector
    vec2search = append(vec2search, entity.FloatVector(embeddingList[0]))
    sRet, err := c.Search(ctx, milvusclient.NewSearchOption(collectionName, topK, vec2search).
        WithPartitions(partitionName).WithOutputFields(idCol, randomCol).WithANNSField(embeddingCol))
    if err != nil {
        t.Fatalf("failed to search collection, err: %v\n", err)
    }
    ids, random, scores := GetResultField(t, sRet[0])
    if len(ids) != topK || len(random) != topK || len(scores) != topK {
        t.Fatalf("len of ret not equal %d vs %d vs %d\n", len(ids), len(random), len(scores))
    }
    t.Logf("search ret ids: %v\n", ids)

    // query, verify delete success
    t.Logf(msgFmt, "start quering pks")
    sliceStr := "["
    for i := 0; i < len(ids)-1; i++ {
        sliceStr += strconv.Itoa(int(ids[i])) + ", "
    }
    sliceStr += strconv.Itoa(int(ids[len(ids)-1]))
    sliceStr += "]"
    expr := fmt.Sprintf("%s in %s", idCol, sliceStr)
    t.Logf(msgFmt, fmt.Sprintf("query with expr `%s`", expr))
    idRet, _ := QueryAndRet(t, ctx, c, []string{partitionName}, expr)
    if len(idRet) != len(ids) {
        t.Fatalf("len of query ret id not equal %d vs %d\n", len(idRet), len(ids))
    }
    t.Logf("query id: %v", idRet)

    t.Logf(msgFmt, "start deleting by pks")
    deleteLen := 2
    deleteIds := ids[0:deleteLen]
    t.Logf("delete list: %v\n", deleteIds)
    pks := column.NewColumnInt64(idCol, deleteIds)
    _, err = c.Delete(ctx, milvusclient.NewDeleteOption(collectionName).
        WithInt64IDs(idCol, deleteIds).WithPartition(partitionName))
    if err != nil {
        t.Fatalf("failed to delete by pks, err: %v\n", err)
    }

    var timeOut = 10
    for ; timeOut > 0; timeOut-- {
        t.Logf(msgFmt, "start querying by pks")
        ret, err := c.Get(ctx, milvusclient.NewQueryOption(collectionName).
            WithPartitions(partitionName).
            WithOutputFields(idCol, randomCol).
            WithIDs(pks))
        if err != nil {
            t.Logf("failed to query result, err: %v", err)
        }
        idRet, _ = GetIDAndRandom(t, ret)
        if err != nil {
            t.Logf("failed to query result, err: %v", err)
        }
        if len(idRet) == 0 {
            t.Logf(msgFmt, "not query deleted data")
            break
        }
        time.Sleep(time.Duration(1) * time.Second)
    }
    if timeOut == 0 {
        t.Fatal("wait Query time out")
    }

    t.Logf(msgFmt, "start searching after deleting by pks")
    newTopK := topK - deleteLen
    sRet, err = c.Search(ctx, milvusclient.NewSearchOption(collectionName, newTopK, vec2search).
        WithPartitions(partitionName).WithOutputFields(idCol, randomCol).WithANNSField(embeddingCol))
    if err != nil {
        t.Fatalf("failed to search collection, err: %v\n", err)
    }
    newIds, newRandom, newScores := GetResultField(t, sRet[0])
    if len(newIds) != newTopK || len(newRandom) != newTopK || len(newScores) != newTopK {
        t.Fatalf("len of ret not equal %d vs %d vs %d\n", len(newIds), len(newRandom), len(newScores))
    }
    t.Logf("newIds: %v\n", newIds)
    for i := 0; i < newTopK; i++ {
        for j := 0; j < len(deleteIds); j++ {
            if newIds[i] == deleteIds[j] {
                t.Fatalf("id %d is deleted, should not be found\n", newIds[i])
            }
        }
    }

    t.Logf(msgFmt, "start adding more")
    // generate data
    idList, randomList, embeddingList = GenerateData(t, nEntities, nEntities)

    // insert data
    InsertPaged(t, ctx, c, idList, randomList, embeddingList, partitionName)
    flushtask, err = c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }

    Search(t, ctx, c, idList, randomList, embeddingList, []string{})
}

// 创建collection，创建index，insert数据，flush，search；再insert数据，flush，search
func TestMoreInsertSearch(t *testing.T) {
    ctx := context.Background()
    t.Logf(msgFmt, "start connecting to Milvus")
    c, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
        Address: milvusAddr,
    })
    if err != nil {
        t.Fatalf("failed to connect to milvus, err: %v\n", err)
    }
    defer c.Close(ctx)

    // delete collection if exists
    has, err := c.HasCollection(ctx, milvusclient.NewHasCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to check collection exists, err: %v\n", err)
    }
    if has {
        c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
    }

    // create collection
    t.Logf(msgFmt, "create collection `hello_ascend_test`")
    schema := CreateSchema()

    if err := c.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
                                 WithShardNum(1)); err != nil {
        t.Fatalf("create collection failed, err: %v\n", err)
    }
    defer func() {
        t.Logf(msgFmt, "drop collection")
        if err := c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName)); err != nil {
            t.Fatalf("failed to drop collection, err: %v\n", err)
        }
    }()

    // build index
    t.Logf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    idx := index.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        t.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v\n", err)
    }
    indexTask, err := c.CreateIndex(ctx, milvusclient.NewCreateIndexOption(collectionName, embeddingCol, idx))
    if err != nil {
        t.Fatalf("failed to create index, err: %v", err)
    }
    if err = indexTask.Await(ctx); err != nil {
        t.Fatalf("Index building failed: %v", err)
    }

    // generate data
    idList, randomList, embeddingList := GenerateData(t, nEntities, 0)

    // insert data
    InsertPaged(t, ctx, c, idList, randomList, embeddingList, "")
    flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }

    // check segment info
    GetSegmentInfo(t, ctx, c)

    // load collection
    t.Logf(msgFmt, "start loading collection")
    loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        t.Fatalf("loading collection failed: %v", err)
    }
    Search(t, ctx, c, idList, randomList, embeddingList, []string{})

    for i := 0; i < 10; i++ {
        // generate more 1000 data
        moreNum := 1000
        idList, randomList, embeddingList = GenerateData(t, moreNum, i*moreNum)

        // insert data
        InsertPaged(t, ctx, c, idList, randomList, embeddingList, "")
        flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
        if err != nil {
            t.Fatalf("failed to flush data, err: %v", err)
        }
        if err = flushtask.Await(ctx); err != nil {
            t.Fatalf("flush failed: %v", err)
        }
        // check segment info
        GetSegmentInfo(t, ctx, c)

        // load collection
        t.Logf(msgFmt, "start loading collection")
        loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
        if err != nil {
            t.Fatalf("failed to load collection, err: %v", err)
        }
        if err = loadTask.Await(ctx); err != nil {
            t.Fatalf("loading collection failed: %v", err)
        }
        Search(t, ctx, c, idList, randomList, embeddingList, []string{})
    }
}

// 一边查询，一边删除
func TestSearchWhileDeleting(t *testing.T) {
    ctx := context.Background()
    t.Logf(msgFmt, "start connecting to Milvus")
    c, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
        Address: milvusAddr,
    })
    if err != nil {
        t.Fatalf("failed to connect to milvus, err: %v\n", err)
    }
    defer c.Close(ctx)

    // delete collection if exists
    has, err := c.HasCollection(ctx, milvusclient.NewHasCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to check collection exists, err: %v\n", err)
    }
    if has {
        c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName))
    }

    // create collection
    t.Logf(msgFmt, "create collection `hello_ascend_test`")
    schema := CreateSchema()

    if err := c.CreateCollection(ctx, milvusclient.NewCreateCollectionOption(collectionName, schema).
                                 WithShardNum(1)); err != nil {
        t.Fatalf("create collection failed, err: %v\n", err)
    }
    defer func() {
        t.Logf(msgFmt, "drop collection")
        if err := c.DropCollection(ctx, milvusclient.NewDropCollectionOption(collectionName)); err != nil {
            t.Fatalf("failed to drop collection, err: %v\n", err)
        }
    }()

    // build index
    t.Logf(msgFmt, "start creating index ASCEND_FLAT_INT8")
    idx := index.NewIndexAscendFlatInt8(entity.IP)
    if err != nil {
        t.Fatalf("failed to create ASCEND_FLAT_INT8 index, err: %v\n", err)
    }
    indexTask, err := c.CreateIndex(ctx, milvusclient.NewCreateIndexOption(collectionName, embeddingCol, idx))
    if err != nil {
        t.Fatalf("failed to create index, err: %v", err)
    }
    if err = indexTask.Await(ctx); err != nil {
        t.Fatalf("Index building failed: %v", err)
    }

    partitionName := "new_partition_1"
    if err = c.CreatePartition(ctx, milvusclient.NewCreatePartitionOption(collectionName, partitionName)); err != nil {
        t.Fatalf("failed to create partition %s, %v\n", partitionName, err)
    }
    has, err = c.HasPartition(ctx, milvusclient.NewHasPartitionOption(collectionName, partitionName))
    if err != nil {
        t.Fatalf("failed to check partition %s, %v\n", partitionName, err)
    }
    if !has {
        t.Fatalf("partition %s created, but not exist", partitionName)
    }

    // generate data
    idList, randomList, embeddingList := GenerateData(t, nEntities, 0)

    // insert data
    InsertPaged(t, ctx, c, idList, randomList, embeddingList, partitionName)
    flushtask, err := c.Flush(ctx, milvusclient.NewFlushOption(collectionName))
    if err != nil {
        t.Fatalf("failed to flush data, err: %v", err)
    }
    if err = flushtask.Await(ctx); err != nil {
        t.Fatalf("flush failed: %v", err)
    }

    // load collection
    t.Logf(msgFmt, "start loading collection")
    loadTask, err := c.LoadCollection(ctx, milvusclient.NewLoadCollectionOption(collectionName))
    if err != nil {
        t.Fatalf("failed to load collection, err: %v", err)
    }
    if err = loadTask.Await(ctx); err != nil {
        t.Fatalf("loading collection failed: %v", err)
    }

    quitSearch := make(chan int)
    t.Logf(msgFmt, "start searching")
    go func(quit <-chan int) {
        for {
            var vec2search []entity.Vector
            vec2search = append(vec2search, entity.FloatVector(embeddingList[0]))
            sRet, err := c.Search(ctx, milvusclient.NewSearchOption(collectionName, topK, vec2search).
                WithPartitions(partitionName).WithOutputFields(idCol, randomCol).WithANNSField(embeddingCol))
            if err != nil {
                t.Fatalf("failed to search collection, err: %v\n", err)
            }
            ids, random, scores := GetResultField(t, sRet[0])
            if len(ids) != topK || len(random) != topK || len(scores) != topK {
                t.Fatalf("len of ret not equal %d vs %d vs %d\n", len(ids), len(random), len(scores))
                <-quit
            }
            t.Logf("search ret ids: %v\n", ids)
            select {
            case <-quit:
                t.Logf("search finished")
                return
            default:
            }
        }
    }(quitSearch)

    for i := 0; i < len(embeddingList)-topK; i++ {
        deleteIds := idList[i : i+1]
        t.Logf("delete list: %v\n", deleteIds)
        _, err := c.Delete(ctx, milvusclient.NewDeleteOption(collectionName).
            WithInt64IDs(idCol, deleteIds).WithPartition(partitionName))
        if err != nil {
            t.Fatalf("failed to delete by pks, err: %v\n", err)
        }
    }
    quitSearch <- 1
}

func GetSegmentInfo(t *testing.T, ctx context.Context, c *milvusclient.Client) {
    // check segment info
    segments, err := c.GetPersistentSegmentInfo(ctx, milvusclient.NewGetPersistentSegmentInfoOption(collectionName))
    if err != nil {
        t.Fatalf("failed to get segment info, err: %v\n", err)
    }
    t.Logf("len of segment %d\n", len(segments))
    for _, segment := range segments {
        t.Logf("segmentId: %d, NumRows: %d, CollectionID: %d, PartitionID %d, State: %v\n",
            segment.ID, segment.NumRows, segment.CollectionID, segment.ParititionID, segment.State)
    }
}

func CreateSchema() *entity.Schema {
    return &entity.Schema{
        CollectionName: collectionName,
        Description:    "hello_ascend_test is the simplest demo to introduce the APIs",
        AutoID:         false,
        Fields: []*entity.Field{
            {
                Name:       idCol,
                DataType:   entity.FieldTypeInt64,
                PrimaryKey: true,
                AutoID:     false,
            },
            {
                Name:       randomCol,
                DataType:   entity.FieldTypeDouble,
                PrimaryKey: false,
                AutoID:     false,
            },
            {
                Name:     embeddingCol,
                DataType: entity.FieldTypeFloatVector,
                TypeParams: map[string]string{
                    entity.TypeParamDim: fmt.Sprintf("%d", dim),
                },
            },
        },
    }
}

func InsertPaged(t *testing.T, ctx context.Context, c *milvusclient.Client, idList []int64, randomList []float64,
    embeddingList [][]float32, partitionName string) {
    t.Logf(msgFmt, "start inserting random entities")
    total := len(idList)
    for i := 0; i < total; i += insertLimit {
        curNum := Min(insertLimit, total-i)
        idColData := column.NewColumnInt64(idCol, idList[i:i+curNum])
        randomColData := column.NewColumnDouble(randomCol, randomList[i:i+curNum])
        embeddingColData := column.NewColumnFloatVector(embeddingCol, dim, embeddingList[i:i+curNum])

        t.Logf("insert %d:%d / %d\n", i, i+curNum, total)
        opt := milvusclient.NewColumnBasedInsertOption(
            collectionName,
            idColData,
            randomColData,
            embeddingColData,
        ).WithPartition(partitionName)
        if _, err := c.Insert(ctx, opt); err != nil {
            log.Fatalf("failed to insert random data into `hello_ascend_test`, err: %v", err)
        }
    }
}

// generate data
func GenerateData(t *testing.T, total int, startID int) ([]int64, []float64, [][]float32) {
    t.Logf("generate data ntotal: %d, dim: %d\n", total, dim)
    rand.Seed(time.Now().UnixNano())
    idList, randomList := make([]int64, total, total), make([]float64, total, total)
    embeddingList := make([][]float32, total, total)
    var wg sync.WaitGroup
    for i := 0; i < total; i++ {
        wg.Add(1)
        go func(i int) {
            idList[i] = int64(i + startID)
            randomList[i] = rand.Float64()
            embeddingList[i] = make([]float32, dim, dim)
            var sum float32 = 0.0
            for j := 0; j < dim; j++ {
                embeddingList[i][j] = rand.Float32()
                sum += embeddingList[i][j] * embeddingList[i][j]
            }
            sum = float32(math.Sqrt(float64(sum)))
            for j := 0; j < dim; j++ {
                embeddingList[i][j] = embeddingList[i][j] / sum
            }
            wg.Done()
        }(i)
    }
    wg.Wait()

    return idList, randomList, embeddingList
}

func SearchExpr(t *testing.T, ctx context.Context, c *milvusclient.Client, embeddingList [][]float32, partitionName []string) {
    t.Logf(msgFmt, "start searching based on vector similarity")

    var vec2search []entity.Vector
    vec2search = append(vec2search, entity.FloatVector(embeddingList[0]))

    // 全量搜索
    sRet, err := c.Search(ctx, milvusclient.NewSearchOption(collectionName, topK, vec2search).
        WithPartitions(partitionName...).WithOutputFields(idCol, randomCol).WithANNSField(embeddingCol))
    if err != nil {
        t.Fatalf("failed to search collection, err: %v\n", err)
    }
    if len(sRet) != len(vec2search) {
        t.Fatalf("len of ret %d not equal batch size %d\n", len(sRet), len(vec2search))
    }

    // ID过滤
    idx, random, scores := GetResultField(t, sRet[0])
    if len(idx) != topK || len(random) != topK || len(scores) != topK {
        t.Fatalf("len of ret not equal %d vs %d vs %d\n", len(idx), len(random), len(scores))
    }
    t.Logf("idx: %v\n", idx)

    excluded := 2
    idxSelect := idx[excluded:]
    sliceStr := "["
    for i := 0; i < len(idxSelect)-1; i++ {
        sliceStr += strconv.Itoa(int(idxSelect[i])) + ", "
    }
    sliceStr += strconv.Itoa(int(idxSelect[len(idxSelect)-1]))
    sliceStr += "]"
    expr := fmt.Sprintf("%s in %s", idCol, sliceStr)
    t.Logf(msgFmt, fmt.Sprintf("search with expr `%s`", expr))
    newTopK := topK - excluded
    sRet, err = c.Search(ctx, milvusclient.NewSearchOption(collectionName, newTopK, vec2search).
        WithFilter(expr).WithPartitions(partitionName...).
        WithOutputFields(idCol, randomCol).WithANNSField(embeddingCol))
    if err != nil {
        t.Fatalf("failed to search collection, err: %v\n", err)
    }
    newIdx, newRandom, newScores := GetResultField(t, sRet[0])
    if len(newIdx) != newTopK || len(newRandom) != newTopK || len(newScores) != newTopK {
        t.Fatalf("len of ret not equal %d vs %d vs %d\n", len(newIdx), len(newRandom), len(newScores))
    }
    t.Logf("newIdx: %v\n", newIdx)
    for i := 0; i < newTopK; i++ {
        if newIdx[i] != idx[i+excluded] && math.Abs(float64(newScores[i])-float64(scores[i+excluded])) > 0.00001 {
            t.Fatalf("idx not equal %d vs %d\n", newIdx[i], idx[i+excluded])
        }
    }
}

func Search(t *testing.T, ctx context.Context, c *milvusclient.Client, idList []int64, randomList []float64,
    embeddingList [][]float32, partitionName []string) {
    t.Logf(msgFmt, "start searching based on vector similarity")
    batchSize := []int{1, 2, 4, 8, 16, 32, 64, 128}
    t.Log("start searching")
    for i := 0; i < len(batchSize); i++ {
        if batchSize[i] > nEntities {
            continue
        }
        t.Logf("test batchsize: %d\n", batchSize[i])

        var idGT []int64
        var randomGT []float64
        var vec2search []entity.Vector
        for j := 0; j < batchSize[i]; j++ {
            idGT = append(idGT, idList[len(idList)-j-1])
            randomGT = append(randomGT, randomList[len(randomList)-j-1])
            vec2search = append(vec2search, entity.FloatVector(embeddingList[len(embeddingList)-j-1]))
        }

        sRet, err := c.Search(ctx, milvusclient.NewSearchOption(collectionName, topK, vec2search).
            WithPartitions(partitionName...).WithOutputFields(idCol, randomCol).WithANNSField(embeddingCol))
        if err != nil {
            t.Fatalf("failed to search collection, err: %v\n", err)
        }
        if len(sRet) != len(vec2search) {
            t.Fatalf("len of ret %d not equal batch size %d\n", len(sRet), len(vec2search))
        }
        if !CheckSearch(t, sRet, idGT, randomGT, topK) {
            t.Fatal("failed to check search result")
        }
    }
}

func GetIDAndRandom(t *testing.T, ret milvusclient.ResultSet) ([]int64, []float64) {
    for _, d := range ret.Fields {
        t.Logf("name:%s", d.Name())
    }
    queryIDCol := ret.GetColumn(idCol)
    var idRetCol *column.ColumnInt64
    iCol, ok := queryIDCol.(*column.ColumnInt64)
    if ok {
        idRetCol = iCol
    } else {
        t.Fatalf("get %s from result failed\n", idCol)
    }
    id := idRetCol.Data()

    queryRandomCol := ret.GetColumn(randomCol)
    var randRetCol *column.ColumnDouble
    rCol, ok := queryRandomCol.(*column.ColumnDouble)
    if ok {
        randRetCol = rCol
    } else {
        t.Fatalf("get %s from result failed\n", randomCol)
    }
    random := randRetCol.Data()
    return id, random
}

func QueryAndRet(t *testing.T, ctx context.Context, c *milvusclient.Client, partitionName []string,
    expr string) ([]int64, []float64) {
    ret, err := c.Query(ctx, milvusclient.NewQueryOption(collectionName).
        WithPartitions(partitionName...).
        WithFilter(expr).
        WithOutputFields(randomCol, idCol))
    if err != nil {
        t.Fatalf("failed to Query, err: %v\n", err)
    }
    return GetIDAndRandom(t, ret)
}

func Query(t *testing.T, ctx context.Context, c *milvusclient.Client, randomList []float64) {
    t.Logf(msgFmt, "query by ID")
    // 随机选择一部分数据，进行Query，对比random是否一致
    sliceAll := Shuffle(nEntities)[0:(nEntities / 300)]
    sliceStr := "["
    for i := 0; i < len(sliceAll)-1; i++ {
        sliceStr += strconv.Itoa(sliceAll[i]) + ", "
    }
    sliceStr += strconv.Itoa(sliceAll[len(sliceAll)-1])
    sliceStr += "]"
    expr := fmt.Sprintf("%s in %s", idCol, sliceStr)
    t.Logf(msgFmt, fmt.Sprintf("query with expr `%s`", expr))
    ret, err := c.Query(ctx, milvusclient.NewQueryOption(collectionName).
        WithFilter(expr).
        WithOutputFields(randomCol))
    if err != nil {
        t.Fatalf("failed to query result, err: %v", err)
    }
    if !CheckQuery(t, ret, randomList, sliceAll) {
        t.Fatal("failed to check query result")
    }
    t.Log("check query result pass")
}

func CheckQuery(t *testing.T, ret milvusclient.ResultSet, randomList []float64, slice []int) bool {
    queryRandomCol := ret.GetColumn(randomCol)
    var randCol *column.ColumnDouble
    c, ok := queryRandomCol.(*column.ColumnDouble)
    if ok {
        randCol = c
    } else {
        t.Errorf("get %s from result failed\n", randomCol)
        return false
    }

    random := randCol.Data()
    sort.Float64s(random)
    var randomGT []float64
    for _, index := range slice {
        randomGT = append(randomGT, randomList[index])
    }
    sort.Float64s(randomGT)
    if len(random) != len(slice) {
        t.Errorf("len of random not equal, %d vs %d", len(random), len(slice))
        return false
    }
    for i := 0; i < len(random); i++ {
        if math.Abs(random[i]-randomGT[i]) > 0.00001 {
            t.Errorf("random[%d] %f not equal %f.\n", i, random[i], randomGT[i])
            return false
        }
    }
    return true
}

func GetResultField(t *testing.T, ret milvusclient.ResultSet) ([]int64, []float64, []float32) {
    var idxCol *column.ColumnInt64
    var randCol *column.ColumnDouble
    for _, field := range ret.Fields {
        if field.Name() == idCol {
            c, ok := field.(*column.ColumnInt64)
            if ok {
                idxCol = c
            } else {
                t.Errorf("get %s from result failed\n", idCol)
            }
        }
        if field.Name() == randomCol {
            c, ok := field.(*column.ColumnDouble)
            if ok {
                randCol = c
            }
        }
    }

    idx := idxCol.Data()
    random := randCol.Data()
    scores := ret.Scores
    return idx, random, scores
}

func CheckSearch(t *testing.T, sRets []milvusclient.ResultSet, idGT []int64, randomGT []float64, topK int) bool {
    if len(sRets) != len(idGT) {
        t.Errorf("len of result not equal, %d vs %d", len(sRets), len(idGT))
        return false
    }

    for i, sRet := range sRets {
        idx, random, scores := GetResultField(t, sRet)
        if len(idx) != topK {
            t.Errorf("len of id not equal, %d vs %d", len(idx), topK)
            return false
        }
        if len(random) != topK {
            t.Errorf("len of random not equal, %d vs %d", len(random), topK)
            return false
        }
        if len(scores) != topK {
            t.Errorf("len of scores not equal, %d vs %d", len(scores), topK)
            return false
        }
        if idx[0] != idGT[i] {
            t.Errorf("top 1 id %d not equal %d.\n", idx[0], idGT[i])
            return false
        }
        if math.Abs(random[0]-randomGT[i]) > 0.00001 {
            t.Errorf("top 1 random %f not equal %f.\n", random[0], randomGT[i])
            return false
        }

        for j := 1; j < len(scores); j++ {
            if scores[j]-scores[j-1] > 0.00001 {
                t.Errorf("scores[%d]:%f > scores[%d]:%f\n", j, scores[j], j-1, scores[j-1])
                return false
            }
        }
    }
    return true
}

func printResult(sRet *milvusclient.ResultSet) {
    pks := make([]string, 0, sRet.ResultCount)
    scores := make([]float32, 0, sRet.ResultCount)
    idCol, ok := sRet.IDs.(*column.ColumnVarChar)
    if !ok {
        panic("not varchar pk")
    }

    randoms := make([]float64, 0, sRet.ResultCount)
    var randCol *column.ColumnDouble
    for _, field := range sRet.Fields {
        if field.Name() == randomCol {
            c, ok := field.(*column.ColumnDouble)
            if ok {
                randCol = c
            }
        }
    }

    for i := 0; i < sRet.ResultCount; i++ {
        randomVal, err := randCol.Value(i)
        if err != nil {
            log.Fatal(err)
        }
        randoms = append(randoms, randomVal)

        val, err := idCol.Value(i)
        if err != nil {
            log.Fatal(err)
        }
        pks = append(pks, val)
        scores = append(scores, sRet.Scores[i])
    }
    fmt.Printf("pks: %v, \tscores: %v\n, \trandoms: %v\n", pks, scores, randoms)
}
